"""
State definition for the attribution analysis agent.
"""
from typing import TypedDict, Optional, List, Dict, Any, Tuple
from datetime import datetime
import pandas as pd


class AttributionState(TypedDict, total=False):
    """
    State object for the attribution analysis workflow.

    This state is passed between nodes in the LangGraph workflow.
    """
    # Input
    query: str                              # User's natural language query

    # Parsed query information
    target_metric: str                      # The metric to analyze (e.g., "revenue", "conversion_rate")
    time_range: Tuple[datetime, datetime]   # Analysis time range
    baseline_period: Tuple[datetime, datetime]  # Baseline comparison period
    dimensions: List[str]                   # Available dimensions for drill-down

    # Data
    data: Optional[pd.DataFrame]            # Raw data for analysis
    metric_metadata: Dict[str, Any]         # Metadata about the metric (formula, unit, etc.)

    # Analysis results
    anomalies: List[Dict[str, Any]]         # Detected anomalies with details
    drill_down_path: List[Dict[str, Any]]   # History of drill-down steps
    current_drill_level: int                # Current drill-down depth
    max_drill_depth: int                    # Maximum allowed drill-down depth

    # Decomposition
    decomposition: Dict[str, Any]           # Metric decomposition results
    contribution_analysis: Dict[str, float] # Contribution of each component

    # Attribution results
    attribution_result: Dict[str, Any]      # Final attribution analysis result
    root_causes: List[Dict[str, Any]]       # Identified root causes

    # Output
    report: str                             # Final natural language report
    visualizations: List[Dict[str, Any]]    # List of visualization specs

    # Control flow
    next_action: str                        # Next action to take in the workflow
    should_continue_drill: bool             # Whether to continue drilling down
    error: Optional[str]                    # Any error message

    # Configuration
    config: Dict[str, Any]                  # Configuration settings
    llm_responses: List[str]                # History of LLM responses for debugging


class DimensionValue(TypedDict):
    """Represents a dimension value with its contribution."""
    dimension: str
    value: str
    contribution: float
    change_rate: float
    is_significant: bool


class AnomalyDetection(TypedDict):
    """Represents an anomaly detection result."""
    metric: str
    timestamp: datetime
    actual_value: float
    expected_value: float
    deviation: float
    severity: str  # 'low', 'medium', 'high'
    confidence: float
    dimensions: Optional[List[DimensionValue]]


class MetricDecomposition(TypedDict):
    """Represents a metric decomposition."""
    metric: str
    formula: str
    components: Dict[str, Any]
    decomposition_type: str  # 'additive', 'multiplicative', 'ratio'
    component_contributions: Dict[str, float]


class DrillDownStep(TypedDict):
    """Represents a single drill-down step."""
    level: int
    dimension: str
    parent_value: Optional[str]
    top_contributors: List[DimensionValue]
    timestamp: datetime
    reasoning: str
